# -*- coding:utf-8 -*-
import pandas as pd
import numpy as np
import os
from datetime import date
from enum import Enum, unique


@unique
class Types(Enum):
    STYLES = 'styles'
    CONCEPT = 'concept'
    EXPONENT = 'exponent'
    TDX_1IND = 'tdx_1ind'
    TDX_CSRC_IND = 'tdx_csrc_ind'


class ExportClass:
    # 源文件路径
    SOURCE_DIR = ''
    # 输出路径
    OUTPUT_DIR = ''
    # 主文件路径
    MAIN_FILE = ''

    # 分类类型
    TYPE_ = ''

    def __init__(self, SOURCE_DIR=None, OUTPUT_DIR=None, MAIN_FILE=None, TYPE_=None):
        if SOURCE_DIR is None and OUTPUT_DIR is None and MAIN_FILE is None and TYPE_ is None:
            print('请输入参数')
            return False
        if TYPE_ not in Types._value2member_map_:
            print('TYPE_ 参数不正确')
            return False
        self.SOURCE_DIR = SOURCE_DIR
        self.OUTPUT_DIR = OUTPUT_DIR
        self.MAIN_FILE = MAIN_FILE
        self.TYPE_ = TYPE_
        self._set_type()

    def _set_type(self):
        """
        设置相关参数
        """
        if self.TYPE_ == Types.STYLES.value:
            self.column_ = '风格'
            self.out_file_name = '沪深A股_风格'
        elif self.TYPE_ == Types.CONCEPT.value:
            self.column_ = '概念'
            self.out_file_name = '沪深A股_概念'

        elif self.TYPE_ == Types.EXPONENT.value:
            self.column_ = '指数'
            self.out_file_name = '沪深A股_指数'
        elif self.TYPE_ == Types.TDX_1IND.value:
            self.column_ = '一级行业'
            self.out_file_name = '沪深A股_通达信一级行业'
        elif self.TYPE_ == Types.TDX_CSRC_IND.value:
            self.column_ = '证监会行业'
            self.out_file_name = '沪深A股_通达信证监会行业'

    def _get_source_filepath(self):
        """
        返回路径下文件文件路径列表
        """
        source_data = os.walk(self.SOURCE_DIR)
        s_files = []
        for root, _, files in source_data:
            for file in files:
                filepath = os.path.join(root, file)
                s_files.append(filepath)
        return s_files

    def _handel_data(self):
        """
        处理数据返回主文件的DataFrame
        1.将所属风格，概念，指数等拼接在 相应的列名下
        """
        files_ = self._get_source_filepath()

        # 读取主文件数据 设置第二列“代码”列为索引
        main_data = pd.read_csv(
            self.MAIN_FILE, dtype={'代码': str}, encoding='gbk', index_col='代码')

        main_data[self.column_] = ''

        # 循环更新处理主文件信息数据
        # names = ''
        for file in files_:

            names = os.path.basename(file).split('_')[0]

            # 设置第一列代码列为索引
            s_data = pd.read_csv(
                file, dtype={'代码': str}, encoding='gbk', index_col='代码')

            for i in s_data.index:
                # print(names)
                # print(str(i))
                temp = main_data[self.column_].get(i)
                if temp is not None:
                    main_data.loc[i, self.column_] = temp+names+'_'
        # 拼接数字列信息
        main_data['num'] = '0.000'
        # 设置无数据的行显示'--'
        main_data[self.column_] = np.where(
            main_data[self.column_] == '', '--', main_data[self.column_])

        main_data.fillna('--')
        # 返回主文件数据信息
        return main_data

    def save_file(self):
        """
        保存数据
        """
        exprot_file_path = self.OUTPUT_DIR+'/'+self.out_file_name + \
            '_'+date.today().strftime('%Y%m%d')+'.txt'
        main_data = self._handel_data()
        # 重置索引为index,不然代码列是索引显示有问题，不能直接导入通达信
        main_data = main_data.reset_index()
        # 代码列补零操作
        main_data['代码'] = ['%06d' % i for i in main_data['代码']]
        # print(main_data)
        # 不导出表头
        main_data.to_csv(exprot_file_path, columns=[
            '市场标识', '代码', self.column_, 'num'], encoding='gbk', header=None, index=False)
        print('*'*10)
        print(self.column_+'数据转化完毕！')
        print('*'*10)
    def gen_name_tocsv(self):
        """
        根据不同分类生成文件名称的csv文件
        """
        files_=self._get_source_filepath()
        pd_data={self.column_:[]}
        
        for file in files_:
            names = os.path.basename(file).split('_')[0]
            pd_data[self.column_].append(names)
        pdd=pd.DataFrame(pd_data)
        pdd.to_csv(self.OUTPUT_DIR+'/'+self.out_file_name+'names.csv',index=False)

if __name__ == '__main__':

    # 源文件路径
    S_STYLES_DIR = '../source/styles'
    S_CONCEPT_DIR = '../source/concept'
    S_EXPONENT_DIR = '../source/exponent'
    S_TDX_1IND_DIR = '../source/tdx_1ind'
    S_TDX_CSRC_IND_DIR = '../source/tdx_csrc_ind'
    # 输出路径
    O_STYLES_DIR = '../out/styles'
    O_CONCEPT_DIR = '../out/concept'
    O_EXPONENT_DIR = '../out/exponent'
    O_TDX_1IND_DIR = '../out/tdx_1ind'
    O_TDX_CSRC_IND_DIR = '../out/tdx_csrc_ind'

    # 主文件路径
    MAIN_FILE = '../source/沪深A股20210618.csv'

    styles_obj = ExportClass(S_STYLES_DIR, O_STYLES_DIR,
                             MAIN_FILE, TYPE_='styles')
    styles_obj.save_file()
    styles_obj.gen_name_tocsv()

    concept_obj = ExportClass(
        S_CONCEPT_DIR, O_CONCEPT_DIR, MAIN_FILE, TYPE_='concept')
    concept_obj.save_file()
    concept_obj.gen_name_tocsv()

    exponent_obj = ExportClass(
        S_EXPONENT_DIR, O_EXPONENT_DIR, MAIN_FILE, TYPE_='exponent')
    exponent_obj.save_file()
    exponent_obj.gen_name_tocsv()

    tdx_1ind_obj = ExportClass(
        S_TDX_1IND_DIR, O_TDX_1IND_DIR, MAIN_FILE, TYPE_='tdx_1ind')
    tdx_1ind_obj.save_file()
    tdx_1ind_obj.gen_name_tocsv()

    tdx_csrc_ind_obj = ExportClass(
        S_TDX_CSRC_IND_DIR, O_TDX_CSRC_IND_DIR, MAIN_FILE, TYPE_='tdx_csrc_ind')
    tdx_csrc_ind_obj.save_file()
    tdx_csrc_ind_obj.gen_name_tocsv()
